﻿// WebSocket.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include "pch.h"
#include "Session.h"
#include "WebSocket.h"
#pragma comment(lib, "WS2_32.lib")

SessionPool WebSocket::sessionPool;
WebSocket::WebSocket(WORD port) :
	IOPort(NULL),
	socketPort(port),
	fnAcceptEx(NULL),
	fnGetAcceptExSockAddrs(NULL),
	acceptCNT(0),
	sessionCNT(0),
	postedAcceptSockets(NULL),
	listenerSocket(INVALID_SOCKET)
{
	WSADATA wsa;
	WSAStartup(MAKEWORD(2, 2), &wsa);

	SYSTEM_INFO info;
	GetSystemInfo(&info);
	numberOfProcessors = info.dwNumberOfProcessors;


	postedAcceptSockets = new PostedAcceptSocket[numberOfProcessors];

	workThreads = new HANDLE[numberOfProcessors];

	for (int i = 0; i < numberOfProcessors; i++) {
		workThreads[i] = INVALID_HANDLE_VALUE;
		postedAcceptSockets[i].socket = INVALID_SOCKET;
	};
	InitializeCriticalSection(&csLock);
	sessions.clear();
}
BOOL WebSocket::start()
{

	GUID guidAcceptEx = WSAID_ACCEPTEX;
	GUID guidGetAcceptExSockAddrs = WSAID_GETACCEPTEXSOCKADDRS;
	DWORD dwBytes = 0;

	//1. 创建一个空的IO端口
	IOPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, numberOfProcessors);
	if (NULL == IOPort) goto ERR;

	//2. 创建监听用socket
	listenerSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (INVALID_SOCKET == listenerSocket) goto ERR_SOCKET;

	//3. 绑定监听socket到IO端口
	CreateIoCompletionPort((HANDLE)listenerSocket, IOPort, 0, 0);

	//4. 绑定Socket到网络端口
	SOCKADDR_IN svrAddr;
	ZeroMemory(&svrAddr, sizeof(svrAddr));
	svrAddr.sin_family = AF_INET;
	svrAddr.sin_addr.s_addr = htonl(INADDR_ANY);
	svrAddr.sin_port = htons(socketPort);
	if (bind(listenerSocket, (SOCKADDR*)&svrAddr, sizeof(svrAddr)) == SOCKET_ERROR) goto ERR_SOCKET;

	//5. 开始监听
	if (listen(listenerSocket, SOMAXCONN) == SOCKET_ERROR) goto ERR_SOCKET;

	//6. 获取内核函数地址
	if (SOCKET_ERROR == WSAIoctl(
		listenerSocket,
		SIO_GET_EXTENSION_FUNCTION_POINTER,
		&guidAcceptEx,
		sizeof(guidAcceptEx),
		&fnAcceptEx,
		sizeof(fnAcceptEx),
		&dwBytes,
		NULL,
		NULL)) goto ERR_SOCKET;
	if (SOCKET_ERROR == WSAIoctl(
		listenerSocket,
		SIO_GET_EXTENSION_FUNCTION_POINTER,
		&guidGetAcceptExSockAddrs,
		sizeof(guidGetAcceptExSockAddrs),
		&fnGetAcceptExSockAddrs,
		sizeof(fnGetAcceptExSockAddrs),
		&dwBytes,
		NULL,
		NULL)) goto ERR_SOCKET;
	//7. 开始准备AcceptSocket;
	for (int i = 0; i < numberOfProcessors; i++) {
		postAccept(postedAcceptSockets[i]);
	}
	if (acceptCNT == 0) goto ERR_ACCEPT;
	//8. 开启工作者线程
	for (int i = 0; i < numberOfProcessors; i++)
	{
		workThreads[i] = CreateThread(0, 0, (LPTHREAD_START_ROUTINE)WorkerThreadProc, (void*)this, 0, 0);
	}
	return true;
ERR_ACCEPT:
	RELEASE(postedAcceptSockets);
ERR_SOCKET:
	RELEASE_SOCKET(listenerSocket);
	RELEASE_HANDLE(IOPort);
ERR:
	return false;
}
void WebSocket::stop()
{
	if (listenerSocket != INVALID_SOCKET) {
		//发送线程关闭事件。
		for (int i = 0; i < numberOfProcessors; i++) {
			PostQueuedCompletionStatus(IOPort, 0, THREAD_EXIT_CMD, NULL);
		}
		//等待线程退出
		WaitForMultipleObjects(numberOfProcessors, workThreads, true, INFINITE);
		//关闭线程句柄
		for (int i = 0; i < numberOfProcessors; i++)
		{
			RELEASE_HANDLE(workThreads[i]);
			acceptCNT--;
			RELEASE_SOCKET(postedAcceptSockets[i].socket);
		}
		for (vector<Session*>::iterator it = sessions.begin(); it != sessions.end(); it++) {
			(*it)->Reset();
			sessionPool.recyleSession(*it);
			sessionCNT--;
		}
		sessions.clear();
		RELEASE_SOCKET(listenerSocket);
		RELEASE_HANDLE(IOPort);
	}

}
WebSocket::~WebSocket()
{
	stop();;
	DeleteCriticalSection(&csLock);
	RELEASE(workThreads);
	RELEASE(postedAcceptSockets);
	WSACleanup();
}

void WebSocket::postAccept(PostedAcceptSocket& postedAcceptSocket)
{
	DWORD dwBytes = 0;
	ZeroMemory(&postedAcceptSocket.ovl, sizeof(WSAOVERLAPPED));
	postedAcceptSocket.ovl.postedAcceptSocket = &postedAcceptSocket;
	postedAcceptSocket.ovl.ioType = IO_TYPE::ACCEPT;
	postedAcceptSocket.socket = WSASocket(AF_INET, SOCK_STREAM, 0, NULL, 0, WSA_FLAG_OVERLAPPED);
	if (INVALID_SOCKET == postedAcceptSocket.socket) return;

	// 将接收缓冲置为0，让Accept时直接返回，用户可以在OnAccept中判断是否接受连接
	if (false == fnAcceptEx(listenerSocket,
		postedAcceptSocket.socket,
		postedAcceptSocket.buf,
		0,
		sizeof(sockaddr_in) + 16,
		sizeof(sockaddr_in) + 16,
		&dwBytes,
		&postedAcceptSocket.ovl))
	{
		if (WSA_IO_PENDING != WSAGetLastError())
		{
			RELEASE_SOCKET(postedAcceptSocket.socket);
			return;
		}
	}
	InterlockedIncrement(&acceptCNT);
}
void WebSocket::Accept(PostedAcceptSocket& postedAcceptSocket)
{
	SOCKADDR_IN* localAddr;
	SOCKADDR_IN* remoteAddr;
	int dwladdr, dwraddr;
	dwladdr = dwraddr = sizeof(SOCKADDR_IN) + 16;
	InterlockedDecrement(&acceptCNT);
	//1. 复制客户端地址
	fnGetAcceptExSockAddrs(postedAcceptSocket.buf,
		0,
		dwladdr,
		dwraddr,
		(sockaddr**)&localAddr,
		&dwladdr,
		(sockaddr**)&remoteAddr,
		&dwraddr);

	//2. 创建新的Session
	Session& session = getSession();
	session.socket = postedAcceptSocket.socket;
	memcpy_s(&session.clientAddr, sizeof(SOCKADDR_IN), remoteAddr, sizeof(SOCKADDR_IN));

	//2. 提交新的Accept
	postAccept(postedAcceptSocket);
	if (onConnect(session)) {
		//3. 绑定到IO端口
		CreateIoCompletionPort((HANDLE)session.socket, IOPort, 0, 0);
		//4. 设置keeplive
		tcp_keepalive inAlive;
		tcp_keepalive outAlive;
		inAlive.onoff = true;
		inAlive.keepalivetime = 60 * 1000;
		inAlive.keepaliveinterval = 10 * 1000;
		DWORD dwBytes;
		WSAIoctl(session.socket, SIO_KEEPALIVE_VALS, (LPVOID)&inAlive, sizeof(inAlive),
			(LPVOID)&outAlive, sizeof(outAlive), &dwBytes, NULL, NULL);
		//5. 准备接收数据
		session.postRecv();
	}
	else {
		debug("from post Access sessioncnt:%d\r\n", sessionCNT);
		recyleSession(&session);
	}
}

Session& WebSocket::getSession()
{
	Session& session = sessionPool.getSession();
	session.ws = this;
	EnterCriticalSection(&csLock);
	sessionCNT++;
	sessions.push_back(&session);
	LeaveCriticalSection(&csLock);
	return session;
}
void WebSocket::recyleSession(Session* session)
{
	onDisconnect(*session);
	session->Reset();

	EnterCriticalSection(&csLock);
	vector<Session*>::iterator it = find(sessions.begin(), sessions.end(), session);
	sessions.erase(it);
	sessionCNT--;
	LeaveCriticalSection(&csLock);

	sessionPool.recyleSession(session);
}
DWORD __stdcall WebSocket::WorkerThreadProc(WebSocket& ws)
{
	DWORD dwBytes;
	OverLapped* ol;
	ULONG	cmd;
	bool rtn;
	while (1) {
		rtn = GetQueuedCompletionStatus(ws.IOPort, &dwBytes, &cmd, (LPOVERLAPPED*)&ol, INFINITE);
		if (THREAD_EXIT_CMD == cmd) break;
		Session& session = *ol->session;
		IO_TYPE type = ol->ioType;
		if (rtn) {
			if ((0 == dwBytes) && (IO_TYPE::RECV == type || IO_TYPE::SEND == type))
			{
				debug("from Threadd A1 sessioncnt%d\r\n", ws.sessionCNT);
				ws.recyleSession(&session);
				continue;
			}
			switch (type) {
			case IO_TYPE::ACCEPT:
				ws.Accept(*ol->postedAcceptSocket);
				break;
			case IO_TYPE::RECV:
			{
				session.Recv(dwBytes);
			}
			break;
			case IO_TYPE::SEND:
				session.sendBufferPool.free(ol->sendBuf.buf);
				ws.onSended(session);
			}
		}
		else {
			DWORD dwErr = GetLastError();
			if (WAIT_TIMEOUT == dwErr) continue;
			debug("from Thread A2,sessioncnt:%d\r\n", ws.sessionCNT);
			ws.recyleSession(&session); 
		}

	}
	return 0;
}
